#!/usr/bin/env python

#python script to analysis the CFG from ssp tool. 
#Copyright to Xingzhong for SSP project in Stevens Institute of Technology, NJ
#Sep, 2011

import pygraphviz as pgv
import networkx as nx
import sys
sys.path.append('./plib')
import libsql as sql
import libgimple as gim
import operator
import numpy as np
from numpy import linalg as LA

def is_parm(node):#{{{
    if node.startswith('PARM_'):
        return True
    else:
        return False#}}}

def is_temp(node):#{{{
    if node.startswith('T.'):
        return True
    else:
        return False#}}}

def is_cst(node):#{{{
    if node.startswith('CONST_'):
        return True
    else:
        return False#}}}

def max_m(M):
    max = 0
    for m,x in M.iteritems():
        for n,y in x.iteritems():
            if y != float('inf'):
                if y > max:
                    max = y
    return max

def sort_dict(d):#{{{
    sd = sorted(d.iteritems(), key=operator.itemgetter(1), reverse=True)
    return sd#}}}

def set_weight(act):
    if act in ["phi"]:
        return 0.1
    elif act in ["nop_expr", "operator[]", "ssa_name", "indirect_ref",\
            "nop_expr", "component_ref", "addr_expr"]:
        return 2
    elif act in ["mult_expr", "plus_expr", "pointer_plus_expr", \
            "minus_expr", "operator+=", "operator*", "operator=",\
            "operator+", "bit_and_expr", "memcpy", "convert_expr",\
            "trunc_mod_expr", "gr_expj", "rdiv_expr"]:
        return 10
    elif act in ["lt_expr", "gt_expr", "ne_expr"]:
        return 20
    else:
        return 1

class GRAPH:
    def __init__(self, name):
        self.name = name.split('/')[-1]
        self.g = nx.DiGraph(name=self.name)
        self.err = 0

    def lsp(self):
        self.parm = []
        self.real_input = []
        self.real_output = []
        self.real_parm = []
        for node in nx.nodes_iter(self.g):
            if is_parm(node):
                self.parm.append(node)
                if node.startswith('PARM_input_'):
                    self.real_input.append(node)
                elif node.startswith('PARM_output_'):
                    self.real_output.append(node)
                else:
                    self.real_parm.append(node)

    def is_connect(self, v1, v2):
        try:
            ret =nx.shortest_path_length(self.g,source=v1,target=v2,weight=True)
        except nx.exception.NetworkXNoPath:    
            ret = -1
        return ret

    def info(self):
        return nx.info(self.g)

    def maxdis(self):
        dist = nx.floyd_warshall_numpy(self.g, weight='weight')
        cond = np.isfinite(dist)
        d = np.extract(cond, dist)
        return np.amax(d) + 1.0
        #return nx.diameter(nx.connected_component_subgraphs(self.g.to_undirected())[0])

    def io_rec(self):
        
        feature = dict(zip(self.parm[:], self.feature[:]))
        feature = sorted(feature, key=feature.get)
        self.IO = {}
        for f in feature:
            self.IO[f] = 'PARM'
            self.tag_pram(f)

        self.IO[feature[0]] = 'OUTPUT'
        self.tag_output(feature[0])
        self.IO[feature[1]] = 'INPUT'
        self.tag_input(feature[1])
        print "[IO]:", self.IO
    
    def u2c(self, U):
        g = self.g
        rank = nx.pagerank(g, max_iter=400)
        #for x,y in rank.iteritems():
        #    print "[PageRank]:",x,y
        gg = g.subgraph(U)
        C = []
        #self.g = gg
        sub = nx.weakly_connected_component_subgraphs(gg)
        for subg in sub:
            print "[u2c subg]",subg
            temp = {}
            for x in subg.nodes(data=True):
                temp[x[0]] = rank[x[0]]
            C.append( sort_dict(temp)[0][0] )
        print C
        return C

    def u2c_2(self, U):
        C = U[:]
        for u1 in U:
            for u2 in U:
                conn = self.is_connect(u1, u2)
                #print "dist[%s,%s]=%s"%(u1,u2,conn)
                if conn > 0 :
                    try:
                        C.remove(u1)
                    except ValueError:
                        pass
        return C


    def u2c_1(self, U):
        g = self.g
        gg = g.subgraph(U)
        G = {}
        for n in gg.nodes(data=True):
            print n
            bb = n[1]['bb']
            if not G.has_key(bb):
                G[bb] = []
            G[bb].append(n[0])
        GG = []
        for x,y in G.iteritems():
            GG.append(g.subgraph(y))
        C = []
        for ggg in GG:
            print ggg.nodes()
            if nx.is_directed_acyclic_graph(ggg):
                c = nx.topological_sort(ggg)[-1]
            else:
                c = ggg.nodes()[0]
            C.append(c)
        return C
    
    def node_style(self, node, col, style='filled'):
        self.g.node[node]['color'] = col
        self.g.node[node]['style'] = style

    def debug(self, type):
        if type == 'node':
            for x in self.g.nodes(data=True):
                print "[debug node]:",x
        elif type == 'edge':
            for x in self.g.edges(data=True):
                print "[debug edge]:",x

    def u2c_3(self, U):
        return U

    def fw(self):
        M = nx.floyd_warshall(self.g, weight='weight')
        Theta = self.parm
        U = []
        # calculate U
        for v in nx.nodes_iter(self.g):
            sum = 0
            for theta in Theta:
                if M[theta][v] != float('inf'):
                    #print "M[%s][%s]=%s"%(theta, v, M[theta][v])
                    sum = sum + 1
            if sum > 1 and v not in self.parm:
                U.append(v)
        for u in U:
            self.node_style(u, 'green4')
            #self.g.node[u]['color'] = 'green4'
        print "[Theta]:", Theta
        print "[U]:", U

        #print "[M]:", M
        # calculate C
        #C = self.u2c_1(U)
        C = self.u2c(U)
        C = self.u2c_3(U)
        #C = self.u2c_2(U)

        print "[C]:", C
        for c in C:
            self.node_style(c, 'green1')
            #self.g.node[c]['color'] = 'green1'

        self.io_u = U
        self.io_c = C
        self.w = []
        m = 0
        n = 0
        distance = self.maxdis()
        distance = 2*max_m(M)
        print "[maxdist]=",distance
        for c in C:
            n = 0
            temp = []
            for theta in Theta:
                try:
                    dist = nx.shortest_path_length(self.g, source=theta,\
                        target=c, weight = True)
                    temp.append( dist )
                except nx.exception.NetworkXNoPath:
                    temp.append( distance )
                n = n+1
            self.w.append(temp)
            m = m +1
        self.W = np.asanyarray(self.w)
        if self.W.size > 0 :
            print self.W
            sum = np.sum(self.W, axis = 1)
            print sum
            C = [C[np.argmin(sum)]]
            print "[C]",C
            for c in C:
                self.node_style(c, 'red')
            self.feature = []

            for w in self.W.T:
                self.feature.append (LA.norm(w))
            print "[feature]", self.feature
            if len( C ) > 0 :
                self.io_rec()
        

    def disp_table(self):
        ret = []
        parm = self.parm[:]
        w = self.w[:]
        ret.append(parm)
        ret.extend(w)
        self.io_c.insert(0, 'header')
        j=0
        for i in ret:
            i.insert(0, self.io_c[j])
            j = j + 1
        temp = ['norm']
        temp.extend (self.feature)
        ret.append(temp)
        temp = ['I/O']
        for i in self.parm:
            temp.append (self.IO[i])
        ret.append(temp)
        return ret

    def work(self):
        
        for node in nx.nodes_iter(self.g):
            if is_parm(node):
                distance = nx.single_source_shortest_path_length(self.g, node)
                for x,y in distance.iteritems():
                    self.g.node[x]['io'] = self.g.node[x]['io'] + 1
                    print "[Distance from] %s to %s is %d[%d]"% (node, x, y,\
                            self.g.node[x]['io'])
        common = {}
        for node in nx.nodes_iter(self.g):
            io = self.is_io(node)
            if io > 0:
                common[node] = io
        scommon = sort_dict(common)
        for x in scommon:
            print '[Common]', x[0], x[1]
        most_common = scommon[0][0]
        self.g.node[most_common]['color'] = 'red'
        print '[Most Common]', most_common

        depth = {}
        for node in self.parm:
            try:
                depth[node] = nx.shortest_path_length(self.g, source = node,\
                    target = most_common)
            except nx.exception.NetworkXNoPath:
                depth[node] = 0
            self.tag_pram(node)


        sdepth = sort_dict(depth)
        for x in sdepth:
            print '[Depth]', x[0], x[1]
        if len(sdepth) > 0:
            self.tag_input(sdepth[0][0])
        if len(sdepth) > 1:
            self.tag_output(sdepth[1][0])
           
    def is_io(self, x):
        return self.g.node[x]['io']

    def tag_input(self, x):
        self.g.node[x]['tag'] = "INPUT"

    def tag_output(self,x):
        self.g.node[x]['tag'] = "OUTPUT"

    def tag_pram(self,x):
        self.g.node[x]['tag'] = "PRAM"

    def eval(self):
        self.err1 = 0
        self.err2 = 0
        self.err3 = 0
        self.sum1 = 0
        self.sum2 = 0
        self.sum3 = 0
        self.sum = len(self.parm)
        self.report = ""
        for node in self.parm:
            print '[Guess]', node, self.g.node[node]['tag']
            self.report = self.report + '[Guess] %s is %s'%(node,\
                    self.g.node[node]['tag'])

        for node in self.real_input:
            self.sum1 = self.sum1 + 1
            if not self.g.node[node]['tag'] == 'INPUT':
                self.err1 = self.err1 + 1

        for node in self.real_output:
            self.sum2 = self.sum2 + 1
            if not self.g.node[node]['tag'] == 'OUTPUT':
                self.err2 = self.err2 + 1

        for node in self.real_parm:
            self.sum3 = self.sum3 + 1
            if not self.g.node[node]['tag'] == 'PRAM':
                self.err3 = self.err3 + 1

        self.err = self.err1 + self.err2 + self.err3
        self.sum = self.sum1 + self.sum2 + self.sum3
        print '[INPUT  ERROR]', self.err1, self.sum1
        print '[OUTPUT ERROR]', self.err2, self.sum2
        print '[PRAM   ERROR]', self.err3, self.sum3
        print '[ERROR]', self.err, self.sum

            

    def draw(self, name= None):
        if name == None:
            name = "%s[err=%d].png"%(self.name, self.err2)
        A = nx.to_agraph(self.g)
        #A.add_node(self.report, shape='box')
        A.layout(prog = 'dot')
        print "drawing cfg..."
        A.draw(name)

    def get(self, Sql):
        for x in Sql.ls():
            X = gim.GIMPLE(x)
            self.add_node2(X)
        #self.merge_node()
        #self.merge_node()

    def add_node2_1(self, node, bb):
        attr = {}
        attr['bb'] = bb
        attr['io'] = 0
        attr['tag'] = 'ANY'
        if is_cst(node):
            return
        elif is_parm(node):
            attr['color'] = 'lightblue'
            attr['style'] = 'filled'
            attr['label'] = node.strip('PARM_')
        elif is_temp(node):
            attr['color'] = 'lightyellow'
            attr['style'] = 'filled'
            #attr['shape'] = 'point'
        else:
            attr['color'] = 'pink'
            attr['style'] = 'filled'
        if not self.g.has_node(node):
            self.g.add_node(node, attr)

    def add_edge2_1(self,n1,n2,act):
        attr = {}
        attr['label'] = act 
        attr['weight'] = set_weight(act)
        if n1.startswith('CONST_') or n2.startswith('CONST_'):
            return
        #if n1 == n2:
        #    return
        if act == 'NULL':
            return
        if act == 'phi':
            self.g.add_edge(n2,n1,attr)
        self.g.add_edge(n1,n2,attr)

    def add_node2(self, node):
        if node.type not in ['gimple_succ', 'gimple_pred']:
            if node.type == 'gimple_phi':
                node.action = 'phi'
            if node.op0 != None:
                self.add_node2_1(node.op0, node.bb)
                if node.op1 != None:
                    self.add_node2_1(node.op1, node.bb)
                    self.add_edge2_1(node.op1, node.op0, node.action)
                if node.op2 != None:
                    self.add_node2_1(node.op2, node.bb)
                    self.add_edge2_1(node.op2, node.op0, node.action)
           
    def merge_node(self):
        print "[in merge_node()]"
        for x in self.g.edges_iter(data=True):
            u = x[0]
            v = x[1]
            data = x[2]
            if data['label'] == 'phi':
                # copy the entire relationship of the v to u
                if self.g.has_node(u):
                    for out in self.g.out_edges_iter(nbunch=[v], data=True):
                        out_node = out[1]
                        out_data = out[2]
                        self.g.add_edge(u, out_node, out_data)
                    for iin in self.g.in_edges_iter(nbunch=[v], data=True):
                        in_node = iin[0]
                        in_data = iin[2]
                        self.g.add_edge(in_node, u, in_data)
                    if self.g.has_node(v):
                        self.g.remove_node(v)
                print "[merge]",u,v,data


